#include "SaveAsDataset.h"
#include "Utils.h"

int main()
{
    Utils::read_config("../config.json");
    auto files = Utils::traverse_folder(Utils::using_img ? "../dataset/images/"
                                                         : "../dataset/videos/");
    SaveAsDataset sad("../dataset/result/");
    cv::Mat img;
    SaveAsDataset::Armors redarmors, bluearmors;
    auto size = Utils::get_size(files);
    for (int lst_i = 0, i = Utils::last_img;
         Utils::using_img ? i < size : i * Utils::video_speed < size;)
    {
        img = Utils::read_img(files, i);
        sad.run(img, Armor::Color::BLUE, bluearmors, Utils::max_find_armor_blue);
        sad.draw_armors(img, bluearmors, -1);
        sad.run(img, Armor::Color::RED, redarmors, Utils::max_find_armor_red);
        sad.draw_armors(img, redarmors, 1);
        auto nxt_i = sad.label_and_save(i, i < lst_i,
                                        {redarmors, bluearmors},
                                        Utils::save_error_img_only);
        lst_i = i;
        i = nxt_i < 0 ? 0 : nxt_i;
    }
    return 0;
}
